rm(list = ls())
library(reshape2)
library(signal)
library(ggplot2)
library(magrittr)
library(stringr)
devtools::load_all()

# 参数设置
k0 <- seq(0.4,16,0.4)
vlast2 <- vlast1 <- rep(20,length(k0))
beta <- .98
delta <-  .1
theta <- .36
numits <- 250
A1 <- 1.75
A2 <- 0.75
p1 <- 0.8
p2 <- 1-p1

#--------值函数迭代---------
# 图形初始化
picdata <- data.frame(k0 = k0, v = vlast1)

# 值函数迭代numits次。在每个k=(0.01，6.2)上以0作为初值开始迭代
v2 <- k2_opt <- v1 <- k1_opt <- numeric(length(k0))
for (i in 1:numits){
  # 寻找每个k点的最优值函数
  for (j in 1:length(k0)){
    # 优化第一个随机状态
    ans <- optimize(valfun_sto, interval = c(0.41,15.99),At=A1, kt = k0[j], beta = beta, theta = theta,
                    delta = delta, k0 = k0, p1=p1,p2=p2,vlast1 = vlast1,vlast2 = vlast2, maximum = T)
    v1[j] <- ans$objective
    k1_opt[j] <- ans$maximum
    # 优化第二个随机状态
    ans <- optimize(valfun_sto, interval = c(0.41,15.99),At=A2, kt = k0[j], beta = beta, theta = theta,
                    delta = delta, k0 = k0, p1=p1,p2=p2,vlast1 = vlast1,vlast2 = vlast2, maximum = T)
    v2[j] <- ans$objective
    k2_opt[j] <- ans$maximum
  }

  # 每50次迭代存储一下
  if (i %% 50 == 0){
    print(i)
    picdata[,paste('v1_',as.character(i), sep = '')] <- v1
    picdata[,paste('k1_opt',as.character(i), sep = '')] <- k1_opt
    picdata[,paste('v2_',as.character(i), sep = '')] <- v2
    picdata[,paste('k2_opt',as.character(i), sep = '')] <- k2_opt
  }

  # 替换上一次的值函数
  vlast1 <- v1
  vlast2 <- v2
}

# value function
ans <- melt(picdata, id.vars = 'k0')
ggplot(ans[str_detect(ans$variable,'^v1'),], aes(x = k0, y = value, color = variable)) +
  geom_line() + geom_line(data = ans[str_detect(ans$variable,'^v2'),],
                          aes(x = k0, y = value, color = variable), linetype = 2) +
  scale_color_manual(values = c(v1_50 = 'red',v1_100 = 'blue',v1_150 = 'green',
                                v1_200 = 'brown',v1_250 = 'black',
                                v2_50 = 'red',v2_100 = 'blue',v2_150 = 'green',
                                v2_200 = 'brown',v2_250 = 'black')) + theme_bw()
# ggsave('../val_ite_sto.png')

# policy function
ggplot(ans[str_detect(ans$variable,'k[1-2]_opt250'),], aes(x = k0, y = value, linetype = variable)) +
  geom_line() + labs(x = 'kt', y = 'ktp1') + theme_bw()
# ggsave('../policy_sto.png')

# 利用政策函数模拟经济中资本存量的变化
simu_k <- data.frame(t = 1:500, k = NA)
simu_k$k[1] <- 1
for (i in 2:nrow(simu_k)) {
  ans <- runif(1)
  if (ans <= 0.8){
    simu_k$k[i] <- signal::interp1(picdata$k0,picdata$k1_opt250,simu_k$k[i-1],'linear')
  }else {
    simu_k$k[i] <- signal::interp1(picdata$k0,picdata$k2_opt250,simu_k$k[i-1],'linear')
  }
}
ggplot(data = simu_k, aes(x = t, y = k)) + geom_line() + theme_bw()
# ggsave('../simu_k.png')

#--------格点法--------
cc <- seq(0.2,0.3,length.out = length(k0)) %>% matrix(ncol = 1)

prob <- data.frame(t1 = c(A1,A1,A2,A2),t2 = c(A1,A2,A1,A2), p = c(p1*p1,p1*p2,p2*p1,p2*p2))
for (i in 1:50) {
  # 计算期望
  EU <- numeric(length(k0))
  for (At1 in c(A1,A2)){# 第t期概率
    knext <- At1*k0^theta + (1-delta)*k0-cc[,ncol(cc)]
    cnext <- interp1(k0,cc[,ncol(cc)],knext,extrap = T)
    for (At2 in c(A1,A2)) {# 第t+1期概率
      EU <- EU + ((At2*theta*knext^(theta-1)+1-delta)/cnext) * prob$p[prob$t1 == At1 & prob$t2 == At2]
    }
  }
  cc <- cbind(cc,1/(beta*EU))
}

picdata <- data.frame(k = k0, cc = cc[,ncol(cc)])
picdata$k175 <- 1.75*k0^theta + (1-delta)*k0-cc[,ncol(cc)]
picdata$k075 <- 0.75*k0^theta + (1-delta)*k0-cc[,ncol(cc)]

# A1,A2两个政策函数
ggplot(picdata, aes(x = k, y = k175)) + geom_line() +
  geom_line(aes(y = k075), linetype = 2) + theme_bw()

# 利用政策函数模拟经济中资本存量的变化
simu_k <- data.frame(t = 1:500, k = NA)
simu_k$k[1] <- 2
for (i in 1:(nrow(simu_k)-1)) {
  cinsert <- interp1(picdata$k,picdata$cc,simu_k$k[i],extrap = T)
  ans <- runif(1)
  if (ans <= 0.8){
    simu_k$k[i+1] <- A1*simu_k$k[i]^theta + (1-delta)*simu_k$k[i]-cinsert
  }else {
    simu_k$k[i+1] <- A2*simu_k$k[i]^theta + (1-delta)*simu_k$k[i]-cinsert
  }
}
ggplot(data = simu_k, aes(x = t, y = k)) + geom_line() + theme_bw()
